/*
 * Copyright 2021 DataCanvas
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.dingodb.calcite.visitor;

import io.dingodb.calcite.rel.DingoAggregate;
import io.dingodb.calcite.rel.DingoDiskAnnBuild;
import io.dingodb.calcite.rel.DingoDiskAnnCountMemory;
import io.dingodb.calcite.rel.DingoDiskAnnLoad;
import io.dingodb.calcite.rel.DingoDiskAnnReset;
import io.dingodb.calcite.rel.DingoDiskAnnStatus;
import io.dingodb.calcite.rel.DingoDocument;
import io.dingodb.calcite.rel.DingoExportData;
import io.dingodb.calcite.rel.DingoFilter;
import io.dingodb.calcite.rel.DingoForUpdate;
import io.dingodb.calcite.rel.DingoFunctionScan;
import io.dingodb.calcite.rel.DingoGenerateSeries;
import io.dingodb.calcite.rel.DingoGetByIndex;
import io.dingodb.calcite.rel.DingoGetByIndexMerge;
import io.dingodb.calcite.rel.DingoGetByKeys;
import io.dingodb.calcite.rel.DingoGetDocumentPreFilter;
import io.dingodb.calcite.rel.DingoGetVectorByDistance;
import io.dingodb.calcite.rel.DingoHybridSearch;
import io.dingodb.calcite.rel.DingoInfoSchemaScan;
import io.dingodb.calcite.rel.DingoLikeScan;
import io.dingodb.calcite.rel.DingoPartCountDelete;
import io.dingodb.calcite.rel.DingoPartRangeDelete;
import io.dingodb.calcite.rel.DingoProject;
import io.dingodb.calcite.rel.DingoReduce;
import io.dingodb.calcite.rel.DingoRepeatUnion;
import io.dingodb.calcite.rel.DingoTableModify;
import io.dingodb.calcite.rel.DingoTableScan;
import io.dingodb.calcite.rel.DingoTableSpool;
import io.dingodb.calcite.rel.DingoTransientTableScan;
import io.dingodb.calcite.rel.DingoUnion;
import io.dingodb.calcite.rel.DingoValues;
import io.dingodb.calcite.rel.DingoVector;
import io.dingodb.calcite.rel.DingoWindow;
import io.dingodb.calcite.rel.DocumentStreamConvertor;
import io.dingodb.calcite.rel.VectorStreamConvertor;
import io.dingodb.calcite.rel.dingo.DingoDocumentScanFilter;
import io.dingodb.calcite.rel.dingo.DingoHashJoin;
import io.dingodb.calcite.rel.dingo.DingoIndexScanWithRelOp;
import io.dingodb.calcite.rel.dingo.DingoReduceAggregate;
import io.dingodb.calcite.rel.dingo.DingoRelOp;
import io.dingodb.calcite.rel.dingo.DingoRoot;
import io.dingodb.calcite.rel.dingo.DingoScanWithRelOp;
import io.dingodb.calcite.rel.dingo.DingoSort;
import io.dingodb.calcite.rel.dingo.DingoStreamingConverter;
import io.dingodb.calcite.rel.dingo.IndexFullScan;
import io.dingodb.calcite.rel.dingo.IndexRangeScan;
import io.dingodb.calcite.visitor.function.DingoAggregateVisitFun;
import io.dingodb.calcite.visitor.function.DingoCountDeleteVisitFun;
import io.dingodb.calcite.visitor.function.DingoDiskAnnBuildVisitFun;
import io.dingodb.calcite.visitor.function.DingoDiskAnnCountMemoryVisitFun;
import io.dingodb.calcite.visitor.function.DingoDiskAnnLoadVisitFun;
import io.dingodb.calcite.visitor.function.DingoDiskAnnResetVisitFun;
import io.dingodb.calcite.visitor.function.DingoDiskAnnStatusVisitFun;
import io.dingodb.calcite.visitor.function.DingoDocumentScanFilterVisitFun;
import io.dingodb.calcite.visitor.function.DingoDocumentStreamingVisitFun;
import io.dingodb.calcite.visitor.function.DingoDocumentVisitFun;
import io.dingodb.calcite.visitor.function.DingoExportDataVisitFun;
import io.dingodb.calcite.visitor.function.DingoFilterVisitFun;
import io.dingodb.calcite.visitor.function.DingoForUpdateVisitFun;
import io.dingodb.calcite.visitor.function.DingoFunctionScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoGenerateSeriesVisitFun;
import io.dingodb.calcite.visitor.function.DingoGetByIndexMergeVisitFun;
import io.dingodb.calcite.visitor.function.DingoGetByIndexVisitFun;
import io.dingodb.calcite.visitor.function.DingoGetByKeysFun;
import io.dingodb.calcite.visitor.function.DingoGetDocumentPreFilterVisitFun;
import io.dingodb.calcite.visitor.function.DingoGetVectorByDistanceVisitFun;
import io.dingodb.calcite.visitor.function.DingoHashJoinVisitFun;
import io.dingodb.calcite.visitor.function.DingoIndexFullScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoIndexRangeScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoIndexScanWithRelOpVisitFun;
import io.dingodb.calcite.visitor.function.DingoInfoSchemaScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoInsertIgnoreVisitFun;
import io.dingodb.calcite.visitor.function.DingoLikeScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoProjectVisitFun;
import io.dingodb.calcite.visitor.function.DingoRangeDeleteVisitFun;
import io.dingodb.calcite.visitor.function.DingoReduceAggregateVisitFun;
import io.dingodb.calcite.visitor.function.DingoReduceVisitFun;
import io.dingodb.calcite.visitor.function.DingoRelOpVisitFun;
import io.dingodb.calcite.visitor.function.DingoRepeatUnionVisitFun;
import io.dingodb.calcite.visitor.function.DingoReplaceInsertVisitFun;
import io.dingodb.calcite.visitor.function.DingoRootVisitFun;
import io.dingodb.calcite.visitor.function.DingoScanWithRelOpVisitFun;
import io.dingodb.calcite.visitor.function.DingoSortVisitFun;
import io.dingodb.calcite.visitor.function.DingoStreamingConverterVisitFun;
import io.dingodb.calcite.visitor.function.DingoTableModifyVisitFun;
import io.dingodb.calcite.visitor.function.DingoTableScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoTableSpoolVisitFun;
import io.dingodb.calcite.visitor.function.DingoTransientTableScanVisitFun;
import io.dingodb.calcite.visitor.function.DingoUnionVisitFun;
import io.dingodb.calcite.visitor.function.DingoValuesVisitFun;
import io.dingodb.calcite.visitor.function.DingoVectorStreamingVisitFun;
import io.dingodb.calcite.visitor.function.DingoVectorVisitFun;
import io.dingodb.calcite.visitor.function.DingoWindowVisitFun;
import io.dingodb.common.ExecuteVariables;
import io.dingodb.common.Location;
import io.dingodb.common.log.LogUtils;
import io.dingodb.exec.base.IdGenerator;
import io.dingodb.exec.base.Job;
import io.dingodb.exec.base.JobManager;
import io.dingodb.exec.dag.Vertex;
import io.dingodb.exec.impl.IdGeneratorImpl;
import io.dingodb.exec.transaction.base.ITransaction;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.sql.SqlKind;
import org.checkerframework.checker.nullness.qual.NonNull;

import java.util.Collection;

import static io.dingodb.calcite.rel.DingoRel.dingo;

@Slf4j
public class DingoJobVisitor implements DingoRelVisitor<Collection<Vertex>> {
    private final IdGenerator idGenerator;
    private final Location currentLocation;
    @Getter
    private final Job job;
    private final ITransaction transaction;

    @Getter
    private final SqlKind kind;

    @Getter
    @Setter
    private boolean isScan;

    @Getter
    private ExecuteVariables executeVariables;

    @Getter
    private long pointTs;

    @Getter
    private boolean forUpdate;

    @Getter
    private boolean replaceInto;

    @Getter
    private boolean isIgnore;

    @Getter
    private long updateLimit;

    @Getter
    private String user;

    @Getter
    private String host;

    private DingoJobVisitor(Job job, IdGenerator idGenerator, Location currentLocation, ITransaction transaction,
                            SqlKind kind, ExecuteVariables executeVariables, long pointTs, boolean forUpdate,
                            boolean replaceInto, boolean isIgnore, long updateLimit, String user, String host) {
        this.job = job;
        this.idGenerator = idGenerator;
        this.currentLocation = currentLocation;
        this.transaction = transaction;
        this.kind = kind;
        this.executeVariables = executeVariables;
        this.pointTs = pointTs;
        this.forUpdate = forUpdate;
        this.replaceInto = replaceInto;
        this.isIgnore = isIgnore;
        this.updateLimit = updateLimit;
        this.user = user;
        this.host = host;
    }

    public static void renderJob(JobManager jobManager, Job job, RelNode input, Location currentLocation) {
        String user = "root";
        String host = "%";
        renderJob(jobManager, job, input, currentLocation, false, null, null,
            new ExecuteVariables(), user, host);
    }

    public static void renderJob(JobManager jobManager, Job job, RelNode input, Location currentLocation,
                                 boolean checkRoot, ITransaction transaction, SqlKind kind,
                                 ExecuteVariables executeVariables, String user, String host) {
        renderJob(jobManager, job, input, currentLocation, checkRoot, transaction, kind,
            executeVariables, 0, false, false, false, -1, user, host);
    }

    public static void renderJob(JobManager jobManager, Job job, RelNode input, Location currentLocation,
                                 boolean checkRoot, ITransaction transaction, SqlKind kind,
                                 ExecuteVariables executeVariables, long pointTs, boolean forUpdate,
                                 boolean replaceInto, boolean isIgnore, long updateLimit, String user, String host) {
        try {
            IdGenerator idGenerator = new IdGeneratorImpl(job.getJobId().seq);
            DingoJobVisitor visitor = new DingoJobVisitor(
                job, idGenerator, currentLocation, transaction, kind,
                executeVariables, pointTs, forUpdate, replaceInto, isIgnore,
                updateLimit, user, host
            );
            Collection<Vertex> outputs = dingo(input).accept(visitor);
            if (checkRoot && !outputs.isEmpty()) {
                throw new IllegalStateException("There root of plan must be `DingoRoot`.");
            }
            if (transaction != null) {
                transaction.setJob(job);
            }

            LogUtils.debug(log, "job = {}", job);
        } catch (Exception e) {
            jobManager.removeJob(job.getJobId());
            throw new RuntimeException(e);
        }
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoStreamingConverter rel) {
        return DingoStreamingConverterVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoAggregate rel) {
        return DingoAggregateVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoFilter rel) {
        return DingoFilterVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoHashJoin rel) {
        return DingoHashJoinVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoTableModify rel) {
        if (replaceInto) {
            return DingoReplaceInsertVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel,
                forUpdate, true, false, updateLimit);
        } else if (isIgnore) {
            return DingoInsertIgnoreVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel,
                forUpdate, false, true, updateLimit);
        }
        return DingoTableModifyVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel,
            forUpdate, false, false, updateLimit);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoProject rel) {
        return DingoProjectVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoReduce rel) {
        return DingoReduceVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoRoot rel) {
        return DingoRootVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoGetByIndex rel) {
        return DingoGetByIndexVisitFun.visit(job, idGenerator, currentLocation, this, transaction, rel);
    }


    @Override
    public Collection<Vertex> visit(@NonNull DingoGetByKeys rel) {
        return DingoGetByKeysFun.visit(job, idGenerator, currentLocation, this, transaction, rel, forUpdate);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoSort rel) {
        return DingoSortVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoTableScan rel) {
        // current version scan must have range
        return DingoTableScanVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoUnion rel) {
        return DingoUnionVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoValues rel) {
        return DingoValuesVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoPartCountDelete rel) {
        return DingoCountDeleteVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoPartRangeDelete rel) {
        return DingoRangeDeleteVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoLikeScan rel) {
        return DingoLikeScanVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoFunctionScan rel) {
        return DingoFunctionScanVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoVector rel) {
        return DingoVectorVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDocument rel) {
        return DingoDocumentVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoHybridSearch rel) {
        throw new UnsupportedOperationException("Not support dingoHybridSearch relNode");
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoGetDocumentPreFilter rel) {
        return DingoGetDocumentPreFilterVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    public Collection<Vertex> visit(@NonNull DingoGetVectorByDistance rel) {
        return DingoGetVectorByDistanceVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull VectorStreamConvertor rel) {
        return DingoVectorStreamingVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DocumentStreamConvertor rel) {
        return DingoDocumentStreamingVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoGetByIndexMerge rel) {
        return DingoGetByIndexMergeVisitFun.visit(job, idGenerator, currentLocation, this, transaction, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoInfoSchemaScan rel) {
        return DingoInfoSchemaScanVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoExportData rel) {
        return DingoExportDataVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull IndexFullScan rel) {
        return DingoIndexFullScanVisitFun.visit(job, idGenerator, currentLocation, this, transaction, rel);
    }

    public Collection<Vertex> visit(@NonNull IndexRangeScan rel) {
        return DingoIndexRangeScanVisitFun.visit(job, idGenerator, currentLocation, this, transaction, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoForUpdate rel) {
        return DingoForUpdateVisitFun.visit(job, idGenerator, currentLocation, this, transaction, rel, forUpdate);
    }

    @Override
    public Collection<Vertex> visitDingoRelOp(@NonNull DingoRelOp rel) {
        return DingoRelOpVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visitDingoScanWithRelOp(@NonNull DingoScanWithRelOp rel) {
        return DingoScanWithRelOpVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visitDingoAggregateReduce(@NonNull DingoReduceAggregate rel) {
        return DingoReduceAggregateVisitFun.visit(job, idGenerator, currentLocation, this, rel);
    }

    @Override
    public Collection<Vertex> visitDingoIndexScanWithRelOp(@NonNull DingoIndexScanWithRelOp rel) {
        return DingoIndexScanWithRelOpVisitFun.visit(job, idGenerator, currentLocation, transaction, this, rel);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDiskAnnStatus dingoDiskAnnStatus) {
        return DingoDiskAnnStatusVisitFun.visit(job, idGenerator, currentLocation, transaction, this, dingoDiskAnnStatus);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDiskAnnCountMemory dingoDiskAnnCountMemory) {
        return DingoDiskAnnCountMemoryVisitFun.visit(job, idGenerator, currentLocation, transaction, this, dingoDiskAnnCountMemory);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDiskAnnReset dingoDiskAnnReset) {
        return DingoDiskAnnResetVisitFun.visit(job, idGenerator, currentLocation, transaction, this, dingoDiskAnnReset);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDiskAnnBuild dingoDiskAnnBuild) {
        return DingoDiskAnnBuildVisitFun.visit(job, idGenerator, currentLocation, transaction, this, dingoDiskAnnBuild);
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoDiskAnnLoad dingoDiskAnnLoad) {
        return DingoDiskAnnLoadVisitFun.visit(job, idGenerator, currentLocation, transaction, this, dingoDiskAnnLoad);
    }

    public Collection<Vertex> visit(@NonNull DingoDocumentScanFilter dingoDocumentScanFilter) {
        return DingoDocumentScanFilterVisitFun.visit(job, idGenerator, currentLocation, this, transaction, dingoDocumentScanFilter);
    }

    public Collection<Vertex> visit(@NonNull DingoWindow dingoWindow) {
        return DingoWindowVisitFun.visit(job, idGenerator, currentLocation, this, transaction, dingoWindow);
    }

    public Collection<Vertex> visit(@NonNull DingoRepeatUnion dingoRepeatUnion) {
        return DingoRepeatUnionVisitFun.visit(
            job, idGenerator, currentLocation, this, transaction, dingoRepeatUnion, executeVariables
        );
    }

    @Override
    public Collection<Vertex> visit(@NonNull DingoGenerateSeries dingoGenerateSeries) {
        return DingoGenerateSeriesVisitFun.visit(
            job, idGenerator, currentLocation, transaction, this, dingoGenerateSeries);
    }

    public Collection<Vertex> visit(@NonNull DingoTransientTableScan dingoTransientTableScan) {
        return DingoTransientTableScanVisitFun.visit(job, idGenerator, currentLocation, this, transaction, dingoTransientTableScan);
    }

    public Collection<Vertex> visit(DingoTableSpool dingoTableSpool) {
        return DingoTableSpoolVisitFun.visit(job, idGenerator, currentLocation, this, transaction, dingoTableSpool);
    }

}
